from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import _init_paths

import os
import cv2
import time

# from opts import opts
from opts_pose import opts
from detectors.detector_factory import detector_factory

image_ext = ['jpg', 'jpeg', 'png', 'webp']
video_ext = ['mp4', 'mov', 'avi', 'mkv']
time_stats = ['tot', 'load', 'pre', 'net', 'dec', 'post', 'merge']

def demo(opt):
  os.environ['CUDA_VISIBLE_DEVICES'] = '0'
  opt.task = 'multi_pose'
  opt.load_model = '/root/project/CenterFace.pytorch/exp/multi_pose/mobilev2_10/model_last.pth' #model_15   model_116_140
  opt.debug = -1
  opt.vis_thresh = 0.3
  Detector = detector_factory[opt.task] # <class 'detectors.multi_pose.MultiPoseDetector'>  opt.task   multi_pose
  detector = Detector(opt)

  imgin = "/root/project/CenterFace.pytorch/src/demo_img_in"
  imgout = "/root/project/CenterFace.pytorch/src/demo_img_out"

  image_names = []
  ls = os.listdir(imgin)
  for file_name in sorted(ls):
      ext = file_name[file_name.rfind('.') + 1:].lower()
      if ext in image_ext:
          image_names.append(os.path.join(imgin, file_name))

  ind = 1
  for (image_name) in image_names:
    # print("image_name",image_name)
    _, outname = os.path.split(image_name)
    outname = 'out_'+outname
    ret = detector.run(image_name)
    # print("detector.run (onedemo.py) success")
    time_str = ''
    for stat in time_stats:
      time_str = time_str + '{} {:.3f}s |'.format(stat, ret[stat])
      # print(ret)
    saveimg = os.path.join(imgout , outname)
    im = ret['plot_img']
    if im.any():
      # print(type(ret['plot_img']),im.shape)
      result = ret['results']
      for bbox in result[1]:
        if bbox[4] > opt.vis_thresh:
          cv2.imwrite(saveimg, im)
          # print(bbox[0:4])
          # print("len of 10",len(bbox[5:]))
          # print("___",outname)
      print(time_str)
    else:
      print("empty")
    ind += 1

def show_results(self, debugger, image, results):
  debugger.add_img(image, img_id='multi_pose')
  for bbox in results[1]:
    if bbox[4] > self.opt.vis_thresh:
      debugger.add_coco_bbox(bbox[:4], 0, bbox[4], img_id='multi_pose')
      debugger.add_coco_hp(bbox[5:39], img_id='multi_pose')
  debugger.show_all_imgs(pause=self.pause)
if __name__ == '__main__':
  opt = opts().init()
  demo(opt)
